package io.gatling.extensions.mongodb.action

import com.typesafe.scalalogging.StrictLogging
import io.gatling.commons.stats.{KO, OK}
import io.gatling.commons.util.Clock
import io.gatling.commons.validation.Failure
import io.gatling.core.CoreComponents
import io.gatling.core.action.{Action, ExitableAction}
import io.gatling.core.check.Check
import io.gatling.core.check.Check.PreparedCache
import io.gatling.core.session.{Expression, Session}
import io.gatling.core.stats.StatsEngine
import io.gatling.core.util.NameGen
import io.gatling.extensions.mongodb.builder._
import io.gatling.extensions.mongodb.protocol.MongoDBComponents
import org.mongodb.scala.bson.conversions.Bson
import org.mongodb.scala.result._
import org.mongodb.scala.{SingleObservable, _}
import org.reactivestreams
import org.reactivestreams.Subscriber

case class MongoDBAction[R](
                             name: String,
                             mongoDBOperate: MongoDBOperate,
                             mongoDBComponents: MongoDBComponents,
                             coreComponents: CoreComponents,
                             throttled: Boolean,
                             next: Action,
                             checkList: List[Check[R]]
                           ) extends ExitableAction with NameGen with StrictLogging {

  import coreComponents._

  override def statsEngine: StatsEngine = coreComponents.statsEngine

  override def clock: Clock = coreComponents.clock

  // todo 看下这里是做什么的
  val preparedCache: PreparedCache = Check.newPreparedCache

  override def execute(session: Session): Unit = {

    mongoDBOperate match {
      case FindOperate(dbName, collectionName, filter) => find(session, dbName, collectionName, filter)
      case InsertOneOperate(dbName, collectionName, document) => insertOneOperate(session, dbName, collectionName, document)
      case InsertManyOperate(dbName, collectionName, documents) => insertManyOperate(session, dbName, collectionName, documents)
      case UpdateOneOperate(dbName, collectionName, filter, update) => updateOneOperate(session, dbName, collectionName, filter, update)
      case UpdateManyOperate(dbName, collectionName, filter, update) => updateManyOperate(session, dbName, collectionName, filter, update)
      case DeleteOneOperate(dbName, collectionName, filter) => deleteOneOperate(session, dbName, collectionName, filter)
      case DeleteManyOperate(dbName, collectionName, filter) => deleteManyOperate(session, dbName, collectionName, filter)
    }
  }

  def find(session: Session, dbName: String, collectionName: String, filter: Expression[String]): Unit = {
    val start: Long = clock.nowMillis
    val result: FindObservable[Document] = mongoDBComponents.client.getDatabase(dbName).getCollection(collectionName).find(Document(filter(session).toOption.get))
    resultOperate(start, session, result)
  }

  def insertOneOperate(session: Session, dbName: String, collectionName: String, document: Expression[String]): Unit = {
    val start: Long = clock.nowMillis
    val result: SingleObservable[InsertOneResult] = mongoDBComponents.client.getDatabase(dbName).getCollection(collectionName).insertOne(Document(document(session).toOption.get))
    resultOperate(start, session, result)
  }

  def insertManyOperate(session: Session, dbName: String, collectionName: String, documents: Seq[_ <: Document]): Unit = {
    val start: Long = clock.nowMillis
    val result: SingleObservable[InsertManyResult] = mongoDBComponents.client.getDatabase(dbName).getCollection(collectionName).insertMany(documents)
    resultOperate(start, session, result)
  }

  def updateOneOperate(session: Session, dbName: String, collectionName: String, filter: Expression[String], update: Expression[String]): Unit = {
    val start: Long = clock.nowMillis
    val result: SingleObservable[UpdateResult] = mongoDBComponents.client.getDatabase(dbName).getCollection(collectionName).updateOne(Document(filter(session).toOption.get), Document(update(session).toOption.get))
    resultOperate(start, session, result)
  }

  def updateManyOperate(session: Session, dbName: String, collectionName: String, filter: Expression[String], update: Seq[Bson]): Unit = {
    val start: Long = clock.nowMillis
    val result: SingleObservable[UpdateResult] = mongoDBComponents.client.getDatabase(dbName).getCollection(collectionName).updateMany(Document(filter(session).toOption.get), update)
    resultOperate(start, session, result)
  }

  def deleteOneOperate(session: Session, dbName: String, collectionName: String, filter: Expression[String]): Unit = {
    val start: Long = clock.nowMillis
    val result: SingleObservable[DeleteResult] = mongoDBComponents.client.getDatabase(dbName).getCollection(collectionName).deleteOne(Document(filter(session).toOption.get))
    resultOperate(start, session, result)
  }

  def deleteManyOperate(session: Session, dbName: String, collectionName: String, filter: Expression[String]): Unit = {
    val start: Long = clock.nowMillis
    val result: SingleObservable[DeleteResult] = mongoDBComponents.client.getDatabase(dbName).getCollection(collectionName).deleteMany(Document(filter(session).toOption.get))
    resultOperate(start, session, result)
  }

  def resultOperate[T](start: Long, session: Session, result: SingleObservable[T]): Unit = {

    result.subscribe(new Subscriber[T] {

      var success: Boolean = true
      val failureInfo: StringBuilder = new StringBuilder

      override def onSubscribe(s: reactivestreams.Subscription): Unit = ???

      override def onNext(r: T): Unit = {
        val checkRes: (Session, Option[Failure]) = Check.check(r.asInstanceOf[R], session, checkList, preparedCache)
        checkRes._2 match {
          case Some(value) => {
            success = false
            failureInfo.append(value.message).append("//")
          }
          case None =>
        }
      }

      override def onError(t: Throwable): Unit = {
        logger.warn("message: {}", t.getMessage)
        statsEngine.logResponse(session.scenario, session.groups, name, start, clock.nowMillis, KO, Some("failed"), Some(t.getMessage))
        throttler match {
          case Some(th) if throttled => th.throttle(session.scenario, () => next ! session)
          case None => next ! session
        }
      }

      override def onComplete(): Unit = {
        if (success) {
          statsEngine.logResponse(session.scenario, session.groups, name, start, clock.nowMillis, OK, None, None)
        } else {
          statsEngine.logResponse(session.scenario, session.groups, name, start, clock.nowMillis, KO, None, Some(failureInfo.toString()))
        }
        throttler match {
          case Some(th) if throttled => th.throttle(session.scenario, () => next ! session)
          case None => next ! session
        }
      }
    })
  }
}
